from keras import backend as K
from tensorflow.keras.layers import Layer
from tensorflow.keras.layers import InputSpec
import tensorflow as tf


def normalize_data_format(value):
    if value is None:
        return 'channels_last'
    data_format = value.lower()
    if data_format not in {'channels_first', 'channels_last'}:
        raise ValueError(f'Invalid data_format: {value}')
    return data_format


def normalize_tuple(value, n, name):
    if isinstance(value, int):
        return (value,) * n
    if isinstance(value, (list, tuple)):
        if len(value) != n:
            raise ValueError(f'The `{name}` argument must be a tuple of {n} integers. Received: {value}')
        return tuple(value)
    raise ValueError(f'The `{name}` argument must be a tuple of {n} integers. Received: {value}')


class BilinearUpsampling(Layer):
    """Just a simple bilinear upsampling layer. Works only with TF.
       Args:
           upsampling: tuple of 2 numbers > 0. The upsampling ratio for h and w
           output_size: used instead of upsampling arg if passed!
    """

    def __init__(self, upsampling=(2, 2), output_size=None, data_format=None, **kwargs):
        super(BilinearUpsampling, self).__init__(**kwargs)

        self.data_format = normalize_data_format(data_format)
        self.input_spec = InputSpec(ndim=4)
        if output_size:
            self.output_size = normalize_tuple(output_size, 2, 'output_size')
            self.upsampling = None
        else:
            self.output_size = None
            self.upsampling = normalize_tuple(upsampling, 2, 'upsampling')

    def compute_output_shape(self, input_shape):
        if self.upsampling:
            height = self.upsampling[0] * input_shape[1] if input_shape[1] is not None else None
            width = self.upsampling[1] * input_shape[2] if input_shape[2] is not None else None
        else:
            height = self.output_size[0]
            width = self.output_size[1]
        return (input_shape[0], height, width, input_shape[3])

    def call(self, inputs):
        if self.upsampling:
            return tf.image.resize(inputs,
                                     (int(inputs.shape[1] * self.upsampling[0]),
                                      int(inputs.shape[2] * self.upsampling[1])),
                                     method='bilinear')
        else:
            return K.tf.image.resize(inputs,
                                     (self.output_size[0], self.output_size[1]),
                                     method='bilinear', align_corners=True)

    def get_config(self):
        config = {'upsampling': self.upsampling,
                  'output_size': self.output_size,
                  'data_format': self.data_format}
        base_config = super(BilinearUpsampling, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
